import { Injectable, Logger } from '@nestjs/common';
import { InjectModel } from '@nestjs/mongoose';
import { Model } from 'mongoose';
import { GenerationResult } from './schemas/generation-result.schema';
import { ComfyUIService } from '../comfyui/comfyui.service';
import { FileUploadService } from '../common/services/file-upload.service';
import { WorkflowTemplateService } from '../workflow-template/workflow-template.service';
import { Workflow } from '../comfyui/interfaces/workflow.interface';

@Injectable()
export class GenerationService {
    private readonly logger = new Logger(GenerationService.name);

    constructor(
        @InjectModel(GenerationResult.name)
        private generationResultModel: Model<GenerationResult>,
        private comfyUIService: ComfyUIService,
        private fileUploadService: FileUploadService,
        private workflowTemplateService: WorkflowTemplateService,
    ) { }

    async generate(templateId: string, params: any, userId?: string) {
        try {
            // 获取模板
            const template = await this.workflowTemplateService.findOne(templateId);

            // 生成图片
            const startTime = Date.now();
            const imageUrl = await this.comfyUIService.generateImage(
                template.workflow as Workflow,
                params
            );

            // 保存生成结果
            const result = await this.generationResultModel.create({
                templateId,
                userId,
                imageUrl,
                params,
                metadata: {
                    generationTime: Date.now() - startTime,
                }
            });

            return result;
        } catch (error) {
            this.logger.error('Generation failed:', error);
            throw error;
        }
    }

    async saveResult(imageUrl: string, templateId: string, params: any, userId?: string) {
        // 下载图片并保存到本地
        const localImageUrl = await this.fileUploadService.downloadAndSaveImage(imageUrl);

        // 创建生成结果记录
        return this.generationResultModel.create({
            templateId,
            userId,
            imageUrl: localImageUrl,
            params,
            metadata: {
                savedAt: new Date(),
            }
        });
    }

    async getResults(userId?: string, page = 1, limit = 10) {
        const query = userId ? { userId } : {};
        const total = await this.generationResultModel.countDocuments(query);
        const results = await this.generationResultModel
            .find(query)
            .sort({ createdAt: -1 })
            .skip((page - 1) * limit)
            .limit(limit)
            .populate('templateId', 'name thumbnail')
            .exec();

        return {
            results,
            pagination: {
                total,
                page,
                limit,
                pages: Math.ceil(total / limit)
            }
        };
    }
} 